import torch
import numpy as np
from termcolor import colored
import pdb

class NumPrint:
    def __init__(self):
        self.print_enable = True
    
    def set_print_enable(self, enable):
        self.print_enable = enable
    
    def print(self, tag, origarr, option="sl", l_dim=-1, force_full_show=False, round_digits=3, override_print=False, todivice=None):
        if origarr is None:
            if self.print_enable:
                print(tag, "is", origarr)
            return
        if not self.print_enable and not override_print:
            return
        arr = origarr
        
        if isinstance(origarr, torch.Tensor):
            arr = origarr.detach().cpu().numpy()
        output = f"{colored(tag, 'green', attrs=['bold'])}"
        if not hasattr(arr, 'shape') or not hasattr(arr, 'ndim'):
            output += f" {arr}"
            print(output)
            return
        if "s" in option:
            output += f"{colored(' shape', 'cyan', attrs=['bold'])}{arr.shape}"
        
        if "l" in option or "c" in option:
            if "l" in option:
                selection = [0] * arr.ndim
                selection[l_dim] = slice(None)
                arr_to_show = arr[tuple(selection)]
            else:
                arr_to_show = arr
            if arr_to_show.flatten().shape[0] > 16 and not force_full_show:
                print(colored(f"{tag}", "green", attrs=['bold']), colored(f"toshowshape {arr_to_show.shape} origshape {arr.shape} first {arr_to_show.flatten()[0]}", "yellow"))
                return
            elif arr_to_show.dtype in [bool, int]:
                output += f" {arr_to_show.astype(int)}"
            all0 = np.all(arr == 0)
            all1 = np.all(arr == 1)
            if all0:
                output += f" [all0:{all0}]"
            elif all1:
                output += f" [all1:{all1}]"
            else:
                rounded_arr = np.round(arr_to_show, round_digits)
                list_arr = list(rounded_arr)
                output += f" {list_arr}"
        if hasattr(arr, 'device'):
            output += f" {arr.device}"
        output += f" {str(type(arr))}"
        print(output)
        if todivice is not None and isinstance(origarr, torch.Tensor):
            return origarr.to(todivice)
        return origarr
p=NumPrint()